import os
import time

from modelscope import AutoModelForCausalLM, AutoTokenizer
import torch
os.environ['MODELSCOPE_CACHE'] = 'D:/hwm_4032442470/Models'#
model_name = "unsloth/Qwen3-8B-unsloth-bnb-4bit"

# 加载 tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 加载 4-bit 量化模型
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_4bit=True,
    device_map="auto",
    torch_dtype=torch.bfloat16
)
print(model.device)
def  qwen3_8B_4bit(prompt):
    # prepare the model input
    #prompt = "Give me a short introduction to large language model."
    messages = [
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False # Switches between thinking and non-thinking modes. Default is True.
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

    # conduct text completion
    generated_ids = model.generate(
        **model_inputs,
        #max_new_tokens=32768,
        max_new_tokens=80,

    )
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

    # parsing thinking content
    try:
        # rindex finding 151668 (</think>)
        index = len(output_ids) - output_ids[::-1].index(151668)
    except ValueError:
        index = 0

    thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

    #print("thinking content:", thinking_content)
    #print("content:", content)
    return content
# 推理示例
if __name__=="__main__":
    startTime=time.time()
    result=qwen3_8B_4bit("有网友发表帖子【六月底能守住3270就行了】他对股市的态度是看涨，看平，看跌，还是分辨不出来态度？，最终必须使用【】将其态度括起来。不要输出我的问题。")
    endTime=time.time()
    print(result)
    print("\n",endTime-startTime)
